import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 读取数据文件
shuju = pd.read_csv("./tools/input.csv", header=0).values
TIME = shuju[:, 0]  # 时间
Rain = shuju[:, 1]  # 时段降雨量
EM = shuju[:, 2]  # 时段蒸发量
QS = shuju[:, 3]  # 时段预报流量
QU = shuju[:, 4]  # 支流汇入，取QU=0

# 输入预报流域面积
Area = 1242  # 水库控制流域面积

# 读取优化好的参数
XX = pd.read_csv("optimized_results.csv", header=0).values.flatten().astype(float)


Sr, Ks, B, alpha0, JL, KKS, KKSS, CS = XX[:8]
z = XX[8:11]
ni = XX[11:14]
H01 = XX[14:17]
Hc1 = XX[17:20]
Hr1 = XX[20:23]
delta = XX[23:26]
LS, LSS, L = map(round, XX[26:29])
Gam = XX[29:32]
EC, KKG, LG = XX[32], XX[33], round(XX[34])

Hs = ni * z
H0 = ni * z * H01
Hc = ni * z * Hc1
Hr = ni * z * Hr1

H = H0.copy()

h0 = H[0]
nz = ni[0] * z[0]
B0 = h0 / nz
dt = TIME[1] - TIME[0]

Nt = len(Rain)
M = len(z)

PE0 = np.zeros(Nt)
G = np.zeros(M)
Gr = np.zeros(M)
E = np.zeros(M)
RSS = np.zeros(Nt)
RS = np.zeros(Nt)
RG = np.zeros(Nt)
RRR = np.zeros(Nt)
balance = np.zeros(Nt)
TRS = np.zeros(Nt)
TRSS = np.zeros(Nt)
TRG = np.zeros(Nt)

Fm, T0, E0 = 0, 0, 0

# 循环计算
for i in range(Nt):
    P = Rain[i]
    # 计算净雨量（扣除蒸发和截留）
    if P >= JL:
        P -= JL
        JL = 0
    else:
        JL -= P
        P = 0
    if P <= EM[i]:
        DEM = EM[i] - P
        PE = 0
        PE0[i] = PE
    else:
        DEM = 0
        PE = P - EM[i]
        PE0[i] = PE

    # 计算地表径流
    B1 = 1 - B0
    Fm1 = Sr * B1 * (T0**0.5) + Ks * T0
    Fm2 = Sr * B1 * ((T0 + dt) ** 0.5) + Ks * (T0 + dt)
    Dfm = Fm2 - Fm1
    X = PE / Dfm

    B1 = B / (1 + B0)
    eta1 = alpha0 + (1 - alpha0) / (B1 + 1)
    if X <= alpha0:
        eta = X
    elif X < 1:
        eta = eta1 - (1 - alpha0) / (B1 + 1) * ((1 - X) / (1 - alpha0)) ** (B1 + 1)
    else:
        eta = eta1
    Df0 = eta * Dfm
    Drb = PE - Df0
    if Drb < 0:
        Drb = 0

    if X <= 1:
        Fm += PE
    else:
        Fm += Dfm

    S0 = Sr * (1 - B0)
    A0 = Ks
    SA = (S0**2 + 4 * A0 * Fm) ** 0.5
    T0 = ((SA - S0) / (2 * A0)) ** 2

    # 壤中流计算
    Drs = 0
    for k in range(M):
        if DEM > 0:
            if H[0] + Df0 >= DEM:
                E[0], E[1], E[2] = DEM, 0, 0
            else:
                E[0] = H[0] + Df0
                if H[1] > EC * Hs[1]:
                    E[1] = (DEM - E[0]) * H[1] / Hs[1]
                    E[2] = 0
                else:
                    if H[1] > EC * (DEM - E[0]):
                        E[1] = EC * (DEM - E[0])
                        E[2] = 0
                    else:
                        E[1] = H[1]
                        E[2] = EC * (DEM - E[0]) - E[1]
        else:
            E[0], E[1], E[2] = 0, 0, 0

        Hw = Df0 + H[k] - E[k]
        beta0 = Hc[k] / Hs[k] / (Gam[k] + 1)
        omega1 = beta0 + (1 - beta0) / (Gam[k] + 1)
        beta = Hw / Hs[k]
        if beta <= beta0:
            omega = beta
        elif beta >= beta0 and beta < 1:
            omega = omega1 - ((1 - beta0) / (Gam[k] + 1)) * (
                (1 - beta) / (1 - beta0)
            ) ** (Gam[k] + 1)
        elif beta >= 1:
            omega = omega1
        Ds = Hw - Hs[k] * omega
        Dr = Ds * delta[k]
        Drs += Dr
        DF = (1 - delta[k]) * Ds
        H[k] = H[k] + Df0 - Ds - E[k]
        Df0 = DF
        if H[k] < 0:
            H[k] = 0
        if H[k] > Hs[k]:
            H[k] = Hs[k]

    RSS[i] = Drs
    RS[i] = Drb
    RG[i] = Df0
    RRR[i] = Drs + Drb + Df0
    E0 += np.sum(E)
    balance[i] = np.sum(H - H0) + np.sum(RRR) - np.sum(PE0) + E0

# 汇流计算
U = Area / dt / 3.6
TRS[:LS] = 0
TRS[LS] = RS[0] * U

for T in range(1 + LS, Nt):
    TRS[T] = TRS[T - 1] * KKS + RS[T - LS] * (1 - KKS) * U

for T in range(1 + LSS, Nt):
    TRSS[T] = TRSS[T - 1] * KKSS + RSS[T - LSS] * (1 - KKSS) * U

TRG[:LG] = QS[0] - QU[0]
for T in range(1 + LG, Nt):
    TRG[T] = TRG[T - 1] * KKG + RG[T - LG] * (1 - KKG) * U

TR = TRS + TRSS + TRG + QU
QJ = TR.copy()
if L < 0:
    L = 0

for T in range(L + 1, Nt):
    QJ[T] = CS * QJ[T - 1] + (1 - CS) * TR[T - L]

# 绘图
plt.plot(TIME, QJ, "b-", linewidth=2)
plt.xlabel("time")
plt.ylabel("flow")
plt.title("predicted flow")
plt.legend(["calculated flow"])
plt.show()

print("peak flow:", max(QJ))
